
#include "stdafx.h"

#include "toolroom.h"
#include "resource.h"

#ifdef _DEBUG
  #define DPRINT DbgPrint
#else
  #define DPRINT
#endif

NTSTATUS  setStatus(PVOID adest, ULONG * palen);
NTSTATUS  setCmpKeyObject();

#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(INIT, setCmpKeyObject)
#pragma alloc_text(PAGE, drvDispatcher)
#pragma alloc_text(PAGE, setStatus)
#pragma alloc_text(PAGE, drvUnload)

PDEVICE_OBJECT deviceObject = NULL;
LPCWSTR strSymlink = L"\\DosDevices\\wygiwys";
LPCWSTR strDrvname = L"\\Device\\wygiwys";
unsigned  modeReadonly = 1;

void callbackThreadNotify(IN HANDLE  ProcessId, 
                          IN HANDLE  ThreadId, 
                          IN BOOLEAN  Create)
{
    removeHandle(ThreadId);
}

//  To disable write protection (WP) bit, that is stored in a special register called control register zero.
__forceinline void CR0DisableMemoryProtection()
{ 
  __asm 
  {
      push eax
      mov eax, CR0
      and eax, 0FFFEFFFFh
      mov CR0, eax
      pop eax
  }    
}

//  To restore write protection (WP) bit.
__forceinline void CR0EnableMemoryProtection()
{
  __asm
  {
      push eax
      mov eax, CR0
      or eax, NOT 0FFFEFFFFh
      mov CR0, eax
      pop eax
  }
}

NTSTATUS setStatus(PVOID adest, ULONG * palen)
{
static LPCWSTR strReadonly= L"READONLY";
static LPCWSTR strWritable= L"WRITABLE";
static const unsigned  replyLen = 9 * sizeof(wchar_t);

    NTSTATUS thestate = STATUS_INVALID_PARAMETER;
    if(palen)
      if(*palen >> 2)
        if(adest)
      {
          ULONG count = replyLen;
          LPCWSTR thource = (modeReadonly) ? strReadonly : strWritable;
          if((*palen) > replyLen) (*palen) = replyLen;
          else {
              RtlZeroMemory(adest, (*palen));
              count = (*palen >> 1) << 1;
              if((*palen) == count) count -= 2;
          }
          RtlCopyMemory(adest, thource, count);
          thestate = STATUS_SUCCESS;
      }
      else *palen = 0;

    return thestate;
}

NTSTATUS drvDispatcher(IN PDEVICE_OBJECT aObject, IN PIRP anIrp)
{
    ULONG thecode = 0;
    ULONG thelen = 0;
    NTSTATUS thestate = STATUS_SUCCESS;
    PIO_STACK_LOCATION thirp = IoGetCurrentIrpStackLocation(anIrp);
    if((!anIrp) || (!thirp))  return STATUS_INVALID_ADDRESS;
  //  ntStatus = STATUS_INVALID_DEVICE_REQUEST;
    if(thirp->MajorFunction == IRP_MJ_DEVICE_CONTROL) 
    {
        thecode = thirp->Parameters.DeviceIoControl.IoControlCode;
        thestate = STATUS_INVALID_PARAMETER;
        thelen = thirp->Parameters.DeviceIoControl.OutputBufferLength;
    };

    DPRINT("_dvk_dbg_drv_, dispatcher by callback, length of offering:\t%u.\n", thelen);
  //if (stack->MajorFunction == IRP_MJ_DEVICE_CONTROL)
    if(thecode)
    __try {
        PVOID buffer = anIrp->AssociatedIrp.SystemBuffer;
        switch(thecode)
        {
        case  IOCTL_QUERY_MODE:
    DPRINT("_dvk_dbg_drv_, dispatcher is query, buffer:\t%P.\n", buffer);
              thestate = setStatus(buffer, &thelen);
          break;

          case  IOCTL_SET_READONLY:
          case  IOCTL_SET_WRITABLE:
#ifdef _DEBUG
              //  WRITABLE is zero
              CR0DisableMemoryProtection();
              __try {
                  modeReadonly = setMode(thecode ^ IOCTL_SET_WRITABLE);
              }
              __finally {
                  CR0EnableMemoryProtection();
              }
    if(modeReadonly)
        DbgPrint("_dvk_dbg_drv_, dispatcher is read only.\n");
    else
        DbgPrint("_dvk_dbg_drv_, dispatcher is WRITABLE.\n");
              thestate = setStatus(buffer, & thelen);
              if(thestate ^ STATUS_SUCCESS)
              {
                  thelen = 0;
                  thestate = STATUS_SUCCESS;
              }
#else
              thestate = setStatus(buffer, & thelen);
#endif
          break;

		      default:
    DPRINT("_dvk_dbg_drv_, dispatcher is STATUS_INVALID_DEVICE_REQUEST.\n");
              thestate = STATUS_INVALID_DEVICE_REQUEST;
        }
    }
    __except(EXCEPTION_EXECUTE_HANDLER)
      {    }

    anIrp->IoStatus.Status = thestate;
    if(thestate ^ STATUS_SUCCESS) thelen = 0;
    anIrp->IoStatus.Information = thelen;

    IoCompleteRequest(anIrp, IO_NO_INCREMENT); 
    return thestate;
}

NTSTATUS drvUnisupported(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	 DPRINT("_dvk_dbg_drv_, happen anything by callback.\n");
   return STATUS_NOT_SUPPORTED;
}

void drvUnload(IN PDRIVER_OBJECT DriverObject)
{
    UNICODE_STRING usDosDeviceName;

    PsRemoveCreateThreadNotifyRoutine(callbackThreadNotify);

    if(deviceObject)
    {
        IoDeleteDevice(deviceObject);
        deviceObject = NULL;
    }

    __try {
        RtlInitUnicodeString(&usDosDeviceName, strSymlink);
        IoDeleteSymbolicLink(&usDosDeviceName);
        CR0DisableMemoryProtection();
        __try {
            restoreSST();  //  restore SST
        }
        __finally {
            CR0EnableMemoryProtection();
        }
    }
    __except(EXCEPTION_EXECUTE_HANDLER)
    {    }

    DPRINT("_dvk_dbg_drv_, unloaded.\n");
}

NTSTATUS DriverEntry(IN PDRIVER_OBJECT aObject, IN PUNICODE_STRING  aRegistryPath)
{
  UNICODE_STRING usDriverName, usDosDevice;
  NTSTATUS thestate = setCmpKeyObject();
  if(NT_ERROR(thestate)) return thestate;
//  if(!(PsIsSystemThread(PsGetCurrentThread()))) return STATUS_INVALID_LEVEL;
	DPRINT("_dvk_dbg_drv_, loaded.\n"); //  \r\n
  {
      OSVERSIONINFOW    theversion;
      theversion.dwOSVersionInfoSize = sizeof(OSVERSIONINFOW);
      thestate = RtlGetVersion(& theversion);
      if(thestate ^ STATUS_SUCCESS) return thestate;
      DPRINT("_dvk_dbg_drv_, os version: %u:%u\n", theversion.dwMajorVersion, theversion.dwMinorVersion);
  //  if(!(theversion.szCSDVersion[0])) return STATUS_UNHANDLED_EXCEPTION;
      if((theversion.dwMajorVersion ^ MAJORVERSION) 
            || (theversion.dwMinorVersion ^ MINORVERSION))
                    return STATUS_REVISION_MISMATCH;  //  os version not fit
  }
  RtlInitUnicodeString(&usDriverName, strDrvname);
  RtlInitUnicodeString(&usDosDevice, strSymlink);

  thestate = IoCreateDevice(aObject, 0,
		                &usDriverName,         // driver name
                    FILE_DEVICE_UNKNOWN,   // Specifying Device Types, FILE_DEVICE_NULL, ???
						        FILE_READ_ONLY_DEVICE, // Specifying Device Characteristics
                    TRUE,                  // If exclusive access to a device is enabled, 
						        &deviceObject);        // only one handle to the device can be open at a time. 
    
  if(thestate == STATUS_SUCCESS)
        thestate = IoCreateSymbolicLink(&usDosDevice, &usDriverName);

  if(thestate == STATUS_SUCCESS)
        thestate = PsSetCreateThreadNotifyRoutine(callbackThreadNotify);

  if(thestate == STATUS_SUCCESS)
  __try {
//    IoCreateDevice( ... FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, 
       unsigned i;
       for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
		        aObject->MajorFunction[i] = drvUnisupported;
//	DriverObject->DriverExtension->AddDevice = biuretAddDevice;
	     aObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = drvDispatcher;
	     aObject->DriverUnload  = drvUnload;
	     aObject->DriverStartIo = NULL;
       CR0DisableMemoryProtection();
       __try {
           thestate = chickyBriky(); //  payload
       }
      __finally {
           CR0EnableMemoryProtection();
      }
  }
  __except(EXCEPTION_EXECUTE_HANDLER)
  {
       thestate = STATUS_DLL_INIT_FAILED;
  };
  
  //  NT_ERROR(thestate)
  if(thestate ^ STATUS_SUCCESS)
  __try {
      drvUnload(NULL);
  }
  __except(EXCEPTION_EXECUTE_HANDLER)
    {    }

	return thestate;
}
